//
// Created by Lenovo on 2021/9/16.
//

#include "cartPoleSystem.h"

MX CartPoleSystem::cartPoleDynamicsMX0(MX x, MX q, MX dx, MX dq, MX F)
{
    return dx;
}
MX CartPoleSystem::cartPoleDynamicsMX1(MX x, MX q, MX dx, MX dq, MX F)
{
    return dq;
}
MX CartPoleSystem::cartPoleDynamicsMX2(MX x, MX q, MX dx, MX dq, MX F)
{
    double T = 0.0;
    MX t2 = cos(q);
    MX t3 = sin(q);
    MX t4 = t2*t2;
    MX t5 = m1+m2-m2*t4;
    MX t6 = 1.0/t5;
    MX t7 = dq*dq;
    double t8 = l*l;
    MX ddx = (t6*(F*l-T*t2+m2*t3*t7*t8+g*l*m2*t2*t3))/l;
    return ddx;
}
MX CartPoleSystem::cartPoleDynamicsMX3(MX x, MX q, MX dx, MX dq, MX F)
{
    double T = 0.0;
    MX t2 = cos(q);
    MX t3 = sin(q);
    MX t4 = t2*t2;
    MX t5 = m1+m2-m2*t4;
    MX t6 = 1.0/t5;
    MX t7 = dq*dq;
    double t8 = l*l;
    double t9 = m2*m2;
    MX ddq = -(1.0/(l*l)*t6*(-T*m1-T*m2+g*l*t3*t9+F*l*m2*t2+g*l*m1*m2*t3+t2*t3*t7*t8*t9))/m2;
    return ddq;
}
double CartPoleSystem::CartPoleInvDynF(double q, double dq, double ddq, double ddx)
{
    double t2 = cos(q);
    double t3 = sin(q);
    return ddx*m1+ddx*m2+ddq*l*m2*t2-dq*dq*l*m2*t3;
}
double CartPoleSystem::CartPoleInvDynT(double q, double dq, double ddq, double ddx)
{
    double t2 = cos(q);
    double t3 = sin(q);
    return l*m2*(ddq*l+ddx*t2+g*t3);
}
GuessStruct CartPoleSystem::computeGuess()
{
    double PI = 3.1415926;
    double tEnd = 2*PI*sqrt(1/g);
    std::vector<double> t,x,dx,ddx,q,dq,ddq,F;
    std::vector<Eigen::Matrix<double, 4,1> > xStates;
    for(int i = 0; i < N; ++i)
    {
        double ti = (tEnd-0.0)/(N-1)*i;
        t.emplace_back(ti);
        x.emplace_back(0.0);
        dx.emplace_back(0.0);
        ddx.emplace_back(0.0);
        q.emplace_back(0.0);
        dq.emplace_back(0.0);
        ddq.emplace_back(0.0);
        F.emplace_back(0.0);
    }
    for(int i = 0; i < N; ++i)
    {
        double xAmp = l*m2/(m1+m2);
        x[i]   =  xAmp*sin(PI*t[i]/tEnd);
        dx[i]  =  xAmp*(2*PI/tEnd)*cos(2*PI*t[i]/tEnd);
        ddx[i] = -xAmp*(2*PI/tEnd)*(2*PI/tEnd)*sin(2*PI*t[i]/tEnd);
        q[i]   = -sign(xAmp)*PI*t[i]/tEnd;
        dq[i]  = -sign(xAmp)*PI*1.0/tEnd;
    }
    for(int i = 0; i < N; ++i)
    {
        Eigen::Matrix<double, 4,1> tmp;
        tmp(0) = x[i];
        tmp(1) = q[i];
        tmp(2) = dx[i];
        tmp(3) = dq[i];
        xStates.emplace_back(tmp);
        F[i] = CartPoleInvDynF(q[i], dq[i], ddq[i], ddx[i]);
    }
    GuessStruct guess;
    guess.time = t;
    guess.state = xStates;
    guess.control = F;
    return guess;
}
BoundsStruct CartPoleSystem::computeBounds(const GuessStruct& guess)
{
    double PI = 3.1415926;
    BoundsStruct bounds;
    bounds.initialState_lower << 0, 0, 0, 0;
    bounds.initialState_upper << 0, 0, 0, 0;
    auto stateEnd = guess.state[guess.state.size()-1];
    bounds.finalState_lower <<stateEnd(0),stateEnd(1), 0, 0;
    bounds.finalState_upper <<stateEnd(0),stateEnd(1), 0, 0;
    double wn = sqrt(1/g);
    bounds.initialTime_lower = 0.0;
    bounds.initialTime_upper = 0.0;
    bounds.finalTime_lower = 0.5*(2*PI*wn);
    bounds.finalTime_upper = 4*(2*PI*wn);
    bounds.state_lower <<-1.0*l, -2*PI, -10*l/wn, -10*PI/wn;
    bounds.state_upper << 1.0*l,  2*PI,  10*l/wn,  10*PI/wn;
    bounds.control_lower = -5*m1*g;
    bounds.control_upper =  5*m1*g;
    return bounds;
}